import torch
import torch.nn as nn
import torch.nn.functional as F

from modules.decoders.base_decoder import BaseDecoder
from modules.multi_head_attn import MultiHeadAttention
from modules.encoders.transformer_encoder import PositionWiseFeedForward
# from onmt.modules.position_ffn import ActivationFunction
from utils.tensor import sequence_mask


class TransformerDecoderLayerBase(nn.Module):
    def __init__(
            self,
            d_model,
            heads,
            d_ff,
            dropout,
            attention_dropout,
            max_relative_positions=0,
            full_context_alignment=False,
            alignment_heads=0,
            pos_ffn_activation_fn=F.relu,
    ):
        """
        Args:
            d_model (int): the dimension of keys/values/queries in
                :class:`MultiHeadedAttention`, also the input size of
                the first-layer of the :class:`PositionwiseFeedForward`.
            heads (int): the number of heads for MultiHeadedAttention.
            d_ff (int): the second-layer of the
                :class:`PositionWiseFeedForward`.
            dropout (float): dropout in residual, self-attn(dot) and
                feed-forward
            attention_dropout (float): dropout in context_attn  (and
                self-attn(avg))
            max_relative_positions (int):
                Max distances between inputs in relative positions
                representations
            full_context_alignment (bool):
                whether enable an extra full context decoder forward for
                alignment
            alignment_heads (int):
                N. of cross attention heads to use for alignment guiding
            pos_ffn_activation_fn (ActivationFunction):
                activation function choice for PositionWiseFeedForward layer

        """
        super(TransformerDecoderLayerBase, self).__init__()

        self.self_attn = MultiHeadAttention(
            head_count=heads,
            model_dim=d_model,
            dropout=attention_dropout,
            max_relative_positions=max_relative_positions,
        )

        self.feed_forward = PositionWiseFeedForward(
            d_model=d_model,
            d_ff=d_ff,
            dropout=0.1,
            activation_fn=pos_ffn_activation_fn)

        self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-6)
        self.drop = nn.Dropout(dropout)
        self.full_context_alignment = full_context_alignment
        self.alignment_heads = alignment_heads

    def forward(self, *args, **kwargs):
        """Extend `_forward` for (possibly) multiple decoder pass:
        Always a default (future masked) decoder forward pass,
        Possibly a second future aware decoder pass for joint learn
        full context alignement, :cite:`garg2019jointly`.

        Args:
            * All arguments of _forward.
            with_align (bool): whether return alignment attention.

        Returns:
            (FloatTensor, FloatTensor, FloatTensor or None):

            * output ``(batch_size, T, model_dim)``
            * top_attn ``(batch_size, T, src_len)``
            * attn_align ``(batch_size, T, src_len)`` or None
        """
        with_align = kwargs.pop("with_align", False)
        output, attns = self._forward(*args, **kwargs)
        top_attn = attns[:, 0, :, :].contiguous()
        attn_align = None
        if with_align:
            if self.full_context_alignment:
                _, attns = self._forward(*args, **kwargs, future=True)

            if self.alignment_heads > 0:
                attns = attns[:, : self.alignment_heads, :, :].contiguous()
            # layer average attention across heads, get ``(B, Q, K)``
            # Case 1: no full_context, no align heads -> layer avg baseline
            # Case 2: no full_context, 1 align heads -> guided align
            # Case 3: full_context, 1 align heads -> full cte guided align
            attn_align = attns.mean(dim=1)
        return output, top_attn, attn_align

    def update_dropout(self, dropout, attention_dropout):
        self.self_attn.update_dropout(attention_dropout)
        self.feed_forward.update_dropout(dropout)
        self.drop.p = dropout

    def _forward(self, *args, **kwargs):
        raise NotImplementedError

    @staticmethod
    def _compute_dec_mask(tgt_pad_mask, future):
        tgt_len = tgt_pad_mask.size(-1)
        if not future:  # apply future_mask, result mask in (B, T, T)
            future_mask = torch.ones(
                [tgt_len, tgt_len],
                device=tgt_pad_mask.device,
                dtype=torch.uint8,
            )
            future_mask = future_mask.triu_(1).view(1, tgt_len, tgt_len)
            # BoolTensor was introduced in pytorch 1.2
            try:
                future_mask = future_mask.bool()
            except AttributeError:
                pass
            dec_mask = torch.gt(tgt_pad_mask + future_mask, 0)
        else:  # only mask padding, result mask in (B, 1, T)
            dec_mask = tgt_pad_mask
        return dec_mask

    def _forward_self_attn(self, inputs_norm, dec_mask, layer_cache, step):
        if isinstance(self.self_attn, MultiHeadAttention):
            return self.self_attn(
                inputs_norm,
                inputs_norm,
                inputs_norm,
                mask=dec_mask,
                layer_cache=layer_cache,
                attn_type="self",
            )
        else:
            raise ValueError(
                f"self attention {type(self.self_attn)} not supported"
            )


class TransformerDecoderLayer(TransformerDecoderLayerBase):
    """Transformer Decoder layer block in Pre-Norm style.
    Pre-Norm style is an improvement w.r.t. Original paper's Post-Norm style,
    providing better converge speed and performance. This is also the actual
    implementation in tensor2tensor and also avalable in fairseq.
    See https://tunz.kr/post/4 and :cite:`DeeperTransformer`.

    .. mermaid::

        graph LR
        %% "*SubLayer" can be self-attn, src-attn or feed forward block
            A(input) --> B[Norm]
            B --> C["*SubLayer"]
            C --> D[Drop]
            D --> E((+))
            A --> E
            E --> F(out)

    """

    def __init__(
            self,
            d_model,
            heads,
            d_ff,
            dropout,
            attention_dropout,
            max_relative_positions=0,
            full_context_alignment=False,
            alignment_heads=0,
            pos_ffn_activation_fn=F.relu,
    ):
        """
        Args:
            See TransformerDecoderLayerBase
        """
        super(TransformerDecoderLayer, self).__init__(
            d_model=d_model,
            heads=heads,
            d_ff=d_ff,
            dropout=dropout,
            attention_dropout=attention_dropout,
            max_relative_positions=max_relative_positions,
            full_context_alignment=full_context_alignment,
            alignment_heads=alignment_heads,
            pos_ffn_activation_fn=pos_ffn_activation_fn,
        )
        self.context_attn = MultiHeadAttention(
            head_count=heads,
            model_dim=d_model,
            dropout=attention_dropout,
            max_relative_positions=0
        )
        self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-6)

    def update_dropout(self, dropout, attention_dropout):
        super(TransformerDecoderLayer, self).update_dropout(
            dropout, attention_dropout
        )
        self.context_attn.update_dropout(attention_dropout)

    def _forward(
            self,
            inputs,
            memory_bank,
            src_pad_mask,
            tgt_pad_mask,
            layer_cache=None,
            step=None,
            future=False,
    ):
        """A naive forward pass for transformer decoder.

        # T: could be 1 in the case of stepwise decoding or tgt_len

        Args:
            inputs (FloatTensor): ``(batch_size, T, model_dim)``
            memory_bank (FloatTensor): ``(batch_size, src_len, model_dim)``
            src_pad_mask : ``(batch_size, 1, src_len)``
            tgt_pad_mask : ``(batch_size, 1, T)``
            layer_cache (dict or None): cached layer info when stepwise decode
            step (int or None): stepwise decoding counter
            future (bool): If set True, do not apply future_mask.

        Returns:
            (FloatTensor, FloatTensor):

            * output ``(batch_size, T, model_dim)``
            * attns ``(batch_size, head, T, src_len)``

        """
        dec_mask = None

        if inputs.size(1) > 1:
            # masking is necessary when sequence length is greater than one
            dec_mask = TransformerDecoderLayerBase._compute_dec_mask(tgt_pad_mask, future)

        inputs_norm = self.layer_norm_1(inputs)
        query, _ = self._forward_self_attn(inputs_norm, dec_mask, layer_cache, step)
        query = self.drop(query) + inputs
        query_norm = self.layer_norm_2(query)
        mid, attns = self.context_attn(
            memory_bank,
            memory_bank,
            query_norm,
            mask=src_pad_mask,
            layer_cache=layer_cache,
            attn_type="context",
        )
        output = self.feed_forward(self.drop(mid) + query)
        return output, attns


if __name__ == '__main__':
    ma = TransformerDecoderLayer(
        d_model=32,
        heads=4,
        d_ff=2,
        dropout=0.,
        attention_dropout=0.,
        max_relative_positions=0,
        full_context_alignment=False,
        alignment_heads=0,
        pos_ffn_activation_fn=F.relu,
    )
    """
                inputs (FloatTensor): ``(batch_size, T, model_dim)``
            memory_bank (FloatTensor): ``(batch_size, src_len, model_dim)``
            src_pad_mask : ``(batch_size, 1, src_len)``
            tgt_pad_mask : ``(batch_size, 1, T)``
            layer_cache (dict or None): cached layer info when stepwise decode
            step (int or None): stepwise decoding counter
            future (bool): If set True, do not apply future_mask.
    """
    t_inputs = torch.zeros(4, 7, 32)
    t_bank = torch.zeros(4, 13, 32)
    tsrc_pad_mask = torch.zeros(4, 1, 13).to(torch.bool)
    ttgt_pad_mask = torch.zeros(4, 1, 7).to(torch.bool)
    t_layer_cache = {}
    t_step = None
    t_future = False
    outputs, attns = ma(
        inputs=t_inputs,
        memory_bank=t_bank,
        src_pad_mask=tsrc_pad_mask,
        tgt_pad_mask=ttgt_pad_mask,
        layer_cache=None,
        step=t_step
    )
